#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
为 documents 集合创建测试数据，解决 knowledge_search 返回空结果的问题
"""

import asyncio
import sys
import os
from pathlib import Path
from datetime import datetime

# 添加项目根目录到路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

# 设置环境变量
os.environ['PYTHONPATH'] = str(project_root / "src")

async def create_knowledge_test_data():
    """为 documents 集合创建测试数据"""
    print("=== 为 documents 集合创建测试数据 ===")
    
    try:
        # 1. 导入必要模块
        print("\n1. 导入模块...")
        
        import chromadb
        from sentence_transformers import SentenceTransformer
        
        # 2. 连接 ChromaDB
        print("\n2. 连接 ChromaDB...")
        
        data_dir = project_root / "data" / "chroma_db"
        chroma_client = chromadb.PersistentClient(path=str(data_dir))
        print(f"✓ ChromaDB 连接成功: {data_dir}")
        
        # 3. 获取或创建 documents 集合
        print("\n3. 获取 documents 集合...")
        
        try:
            documents_collection = chroma_client.get_collection('documents')
            print("✓ documents 集合已存在")
        except Exception:
            # 如果集合不存在，创建它
            documents_collection = chroma_client.create_collection(
                name='documents',
                metadata={"hnsw:space": "cosine"}
            )
            print("✓ documents 集合已创建")
        
        # 4. 检查当前数据
        print("\n4. 检查当前数据...")
        current_count = documents_collection.count()
        print(f"当前 documents 集合文档数量: {current_count}")
        
        # 5. 准备测试文档数据
        print("\n5. 准备测试文档数据...")
        
        test_documents = [
            {
                "id": "doc_mcp_test_001",
                "content": "MCP工具验证测试：这是一个用于测试MCP（Model Context Protocol）工具功能的文档。MCP是一个用于AI模型与外部工具交互的协议标准。",
                "metadata": {
                    "title": "MCP工具验证测试文档",
                    "category": "测试文档",
                    "tags": "MCP,工具,验证,测试",
                    "created_at": datetime.now().isoformat(),
                    "user_id": "default_user",
                    "document_type": "test"
                }
            },
            {
                "id": "doc_mcp_protocol_001",
                "content": "MCP协议介绍：Model Context Protocol (MCP) 是一个开放标准，用于连接AI助手与各种数据源和工具。它提供了统一的接口来访问文件系统、数据库、API等资源。",
                "metadata": {
                    "title": "MCP协议介绍",
                    "category": "技术文档",
                    "tags": "MCP,协议,AI,工具",
                    "created_at": datetime.now().isoformat(),
                    "user_id": "default_user",
                    "document_type": "documentation"
                }
            },
            {
                "id": "doc_vue_framework_001",
                "content": "Vue.js框架介绍：Vue是一个用于构建用户界面的渐进式JavaScript框架。它易于学习，具有响应式数据绑定和组件化开发的特点。",
                "metadata": {
                    "title": "Vue.js框架介绍",
                    "category": "前端开发",
                    "tags": "Vue,JavaScript,前端,框架",
                    "created_at": datetime.now().isoformat(),
                    "user_id": "default_user",
                    "document_type": "documentation"
                }
            },
            {
                "id": "doc_testing_guide_001",
                "content": "软件测试指南：软件测试是确保软件质量的重要环节。包括单元测试、集成测试、系统测试等多个层次。测试驱动开发(TDD)是一种有效的开发方法。",
                "metadata": {
                    "title": "软件测试指南",
                    "category": "软件工程",
                    "tags": "测试,软件工程,TDD,质量保证",
                    "created_at": datetime.now().isoformat(),
                    "user_id": "default_user",
                    "document_type": "guide"
                }
            },
            {
                "id": "doc_python_programming_001",
                "content": "Python编程基础：Python是一种高级编程语言，具有简洁的语法和强大的功能。广泛应用于Web开发、数据科学、人工智能等领域。",
                "metadata": {
                    "title": "Python编程基础",
                    "category": "编程语言",
                    "tags": "Python,编程,开发,语言",
                    "created_at": datetime.now().isoformat(),
                    "user_id": "default_user",
                    "document_type": "tutorial"
                }
            }
        ]
        
        print(f"准备添加 {len(test_documents)} 条测试文档")
        
        # 6. 添加测试数据到 ChromaDB
        print("\n6. 添加测试文档...")
        
        # 提取数据
        ids = [doc["id"] for doc in test_documents]
        documents = [doc["content"] for doc in test_documents]
        metadatas = [doc["metadata"] for doc in test_documents]
        
        # 添加到集合
        documents_collection.add(
            ids=ids,
            documents=documents,
            metadatas=metadatas
        )
        
        print(f"✓ 成功添加 {len(test_documents)} 条测试文档")
        
        # 7. 验证数据添加
        print("\n7. 验证数据添加...")
        new_count = documents_collection.count()
        added_count = new_count - current_count
        print(f"添加后 documents 集合文档数量: {new_count}")
        print(f"新增文档数量: {added_count}")
        
        # 8. 测试搜索功能
        print("\n8. 测试搜索功能...")
        
        # 测试搜索 "MCP工具验证测试"
        search_results = documents_collection.query(
            query_texts=["MCP工具验证测试"],
            n_results=3
        )
        
        print(f"搜索 'MCP工具验证测试': {len(search_results['documents'][0])} 个结果")
        
        if search_results['documents'][0]:
            print("\n搜索结果详情:")
            for i, (doc_id, content, metadata, distance) in enumerate(zip(
                search_results['ids'][0],
                search_results['documents'][0],
                search_results['metadatas'][0],
                search_results['distances'][0]
            )):
                print(f"  结果 {i+1}:")
                print(f"    ID: {doc_id}")
                print(f"    标题: {metadata.get('title', 'N/A')}")
                print(f"    内容: {content[:100]}...")
                print(f"    相似度距离: {distance:.4f}")
                print()
        
        # 9. 验证其他搜索词
        print("\n9. 验证其他搜索词...")
        
        test_queries = ["Vue框架", "Python编程", "软件测试"]
        for query in test_queries:
            results = documents_collection.query(
                query_texts=[query],
                n_results=2
            )
            print(f"搜索 '{query}': {len(results['documents'][0])} 个结果")
        
        print("\n=== 知识库测试数据创建完成 ===")
        print("\n现在可以重新测试 MCP knowledge_search 工具:")
        print('搜索参数: {"query": "MCP工具验证测试"}')
        print("预期结果: 应该能找到相关的知识文档")
        
    except Exception as e:
        print(f"❌ 创建测试数据时出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(create_knowledge_test_data())